import random
import numpy
from Area import Area
from Person import Person
import time

# Meta 方法汇总
"""
run1day(bool 检查点) 世界运行一天 检查点弃用
  返回一天3阶段人员位置分布
getPersonNum()
  总人数
getDays()
  当前天数
getHospitalCapacity()
  医院的总容量，剩余容量
getIsolationCapacity()
  隔离点的总容量，剩余容量
getHospitalCare()
  医院的治愈率，死亡率
getIsolationCare()
  隔离点的治愈率，死亡率
getPersonStatus()
  健康人数，潜伏人数，无症状人数，确诊人数，死亡人数
getPatientCum()
  确诊人次
detection() 核酸检测
  False 医院隔离点剩余空间不足
areasHandler()
  返回每个区域，对应各种状态人的数目
getIncrease()
  返回新增确诊、治愈和死亡
setHospitalRisk(float 医院传染风险)
setLibraryRisk(float 图书馆传染风险)
setIsolationRisk(float 隔离点传染风险)
setOutdoorRisk(float 室外传染风险)
setClassroomRisk(float 教室传染风险)
setDormRisk(float 宿舍传染风险)
setHospitalCureRate(float 医院治愈率)
setHospitalDeadRate(float 医院死亡率)
setIsolationCureRate(float 隔离点治愈率)
setIsolationDeadRate(float 隔离点死亡率)
setPatient0(int 0号病人数量, float 0号病人携带病毒传染率)
vaccinate(float 疫苗有效率，float 接种比例)
allMask(float 口罩有效率)
morgue(Person 死人)
setConfront(bool 是否积极抗疫)

# 非外部调用
move
statusChange
infect
errorCheck
"""

"""
hospital, isolation, outdoor, library 都设置为1个且不改变 但可改变其容量
id = 0, 1, 2, 3
如必要 图书馆可以设置多个
人数预置后不可更改 后期考虑增加人数设置
"""


class Meta:
    """
    _risk_level  dict  场所分类及风险等级
    _area_num  dict  各类场所数量
    _area_capacity  dict  各类场所承载量
    _area_cure  dict  各场景治愈率
    _area_dead dict  各场景死亡率
    _person_num  int  人数
    _total_area  int  总场所数量
    _areas  list  场所
    _hospital_id  int 医院id开始
    _isolation_id  int  隔离所id开始
    _outdoor_id  int  室外id开始
    _library_id  int  图书馆id开始
    _classroom_id  int  教室id开始
    _dorm_id  int  教室id开始
    p_health  list  健康人
    p_hide  list  未病发无感染性
    p_no_sym  list  无症状有传染性
    p_patient  list  病发人
    p_dead  list  死人
    _days  int  天数计时
    _confront  bool  抗击疫情 0无作为 1病发就进医院，医院没位置进隔离所
    _patient_cum  int  确诊人次
    """

    def __init__(self, person_num, hospital_cap, isolation_cap, library_cap, classroom_num):
        # 场所分类及风险等级
        self._area_risk = {'dorm': 0.8, 'classroom': 0.5, 'library': 0.2, 'outdoor': 0, 'hospital': 0, 'isolation': 0}
        # 各类场所数量
        self._area_num = {'dorm': (person_num + 3) // 4, 'classroom': classroom_num, 'library': 1, 'outdoor': 1,
                          'hospital': 1, 'isolation': 1}
        # 各类场所承载量
        self._area_capacity = {'dorm': 4, 'classroom': 100, 'library': library_cap, 'outdoor': 99999,
                               'hospital': hospital_cap, 'isolation': isolation_cap}
        # 各场景的治愈率
        self._area_cure = {'hospital': 0.4, 'isolation': 0.2, 'else': 0.1}
        # 各场景的死亡率
        self._area_dead = {'hospital': 0.006, 'isolation': 0.01, 'else': 0.03}

        # 人数
        self._total_person = person_num
        # 场所计数
        self._total_area = 0
        # 场所
        self._areas = []
        # 健康人
        self.p_health = []
        # 潜伏
        self.p_hide = []
        # 无症状有传染性
        self.p_no_sym = []
        # 病人
        self.p_patient = []
        # 死人
        self.p_dead = []

        # 添加医院
        self._areas.append(
            Area(self._total_area, self._area_risk['hospital'], self._area_capacity['hospital'], 'hospital', 0))
        self._hospital_id = self._total_area
        self._total_area += 1
        # 添加隔离所
        self._areas.append(
            Area(self._total_area, self._area_risk['isolation'], self._area_capacity['isolation'], 'isolation', 1))
        self._isolation_id = self._total_area
        self._total_area += 1
        # 添加室外
        self._areas.append(
            Area(self._total_area, self._area_risk['outdoor'], self._area_capacity['outdoor'], 'outdoor', 2))
        self._outdoor_id = self._total_area
        self._total_area += 1
        # 添加图书馆
        self._areas.append(
            Area(self._total_area, self._area_risk['library'], self._area_capacity['library'], 'library', 3))
        self._library_id = self._total_area
        self._total_area += 1
        # 添加教室
        self._classroom_id = self._total_area
        for aid in range(self._area_num['classroom']):
            self._areas.append(
                Area(self._total_area, self._area_risk['classroom'], self._area_capacity['classroom'], 'classroom', 4))
            self._total_area += 1
        # 添加人和宿舍
        self._dorm_id = self._total_area
        pid = 0
        for aid in range(self._area_num['dorm']):
            self._areas.append(Area(self._total_area, self._area_risk['dorm'], self._area_capacity['dorm'], 'dorm', 5))
            for j in range(4):
                if pid < self._total_person:
                    self.p_health.append(Person(self._total_area))
                    pid += 1
            self._areas[self._total_area].setRemain(0)
            self._total_area += 1

        assert pid == self._total_person
        assert self._total_area == sum(self._area_num.values())

        # 天数
        self._days = 0
        # 是否积极抗击疫情
        self._confront = False
        # 确诊人次
        self._patient_cum = 0
        # 新增确诊治愈和死亡
        self._increase = [0 for i in range(3)]

    def getPersonNum(self):
        return self._total_person

    def getDays(self):
        return self._days

    def getHospitalCapacity(self):
        return self._area_capacity['hospital'], self._areas[self._hospital_id].getRemain()

    def getIsolationCapacity(self):
        return self._area_capacity['isolation'], self._areas[self._isolation_id].getRemain()

    def getHospitalCare(self):
        return self._area_cure['hospital'], self._area_dead['hospital']

    def getIsolationCare(self):
        return self._area_cure['isolation'], self._area_dead['isolation']

    def getPersonStatus(self):
        return len(self.p_health), len(self.p_hide), len(self.p_no_sym), len(self.p_patient), len(
            self.p_dead), self._total_person

    def getPatientCum(self):
        return self._patient_cum

    def getIncrease(self):
        # 分别是 新增确诊 新增治愈 新增死亡
        return self._increase

    def getConfront(self):
        return self._confront

    def setConfront(self, flag):
        if flag:
            self._confront = True
        else:
            self._confront = False

    # 设置医院风险
    def setHospitalRisk(self, rate):
        self._area_risk['hospital'] = rate
        self._areas[self._hospital_id].setRisk(rate)

    # 设置图书馆风险
    def setLibraryRisk(self, rate):
        self._area_risk['library'] = rate
        self._areas[self._library_id].setRisk(rate)

    # 设置隔离点风险
    def setIsolationRisk(self, rate):
        self._area_risk['isolation'] = rate
        self._areas[self._isolation_id].setRisk(rate)

    # 设置室外风险
    def setOutdoorRisk(self, rate):
        self._area_risk['outdoor'] = rate
        self._areas[self._outdoor_id].setRisk(rate)

    # 设置教室风险
    def setClassroomRisk(self, rate):
        self._area_risk['classroom'] = rate
        for ids in range(self._classroom_id, self._classroom_id + self._area_num['classroom']):
            self._areas[ids].setRisk(rate)

    # 设置宿舍风险
    def setDormRisk(self, rate):
        self._area_risk['dorm'] = rate
        for ids in range(self._dorm_id, self._dorm_id + self._area_num['dorm']):
            self._areas[ids].setRisk(rate)

    # 设置医院治愈率
    def setHospitalCureRate(self, rate):
        self._area_cure['hospital'] = rate

    # 设置医院死亡率
    def setHospitalDeadRate(self, rate):
        self._area_dead['hospital'] = rate

    # 设置隔离点治愈率
    def setIsolationCureRate(self, rate):
        self._area_cure['isolation'] = rate

    # 设置隔离点死亡率
    def setIsolationDeadRate(self, rate):
        self._area_dead['isolation'] = rate

    # 随机设置0号病人 设置病人初始为无症状 有传染性 已感染3天 virus病毒感染率
    def setPatient0(self, num=1, virus=0.7):
        num = max(0, int(num))
        virus = max(0.5, min(1.0, virus))
        for i in range(num):
            person = random.choice(self.p_health)
            person.bePatient0(virus)
            self.p_health.remove(person)
            self.p_no_sym.append(person)

    # 人员流动
    def move(self, person, area):
        before = person.moveToArea(area)
        self._areas[before].addRemain()
        self._areas[area].reduceRemain()

    # 人死进太平间
    def morgue(self, person):
        before = person.moveToArea(-1)
        self._areas[before].addRemain()

    # 健康状态变化
    def statusChange(self, person, before, after):
        if before == 0:
            self.p_health.remove(person)
        elif before == 1:
            self.p_hide.remove(person)
        elif before == 2:
            self.p_no_sym.remove(person)
        elif before == 3:
            self.p_patient.remove(person)
        if after == 0:
            self.p_health.append(person)
        elif after == 1:
            self.p_hide.append(person)
        elif after == 2:
            self.p_no_sym.append(person)
        elif after == 3:
            self.p_patient.append(person)
            self._patient_cum += 1
        else:
            self.p_dead.append(person)

    # 一天演进
    def run1day(self, _check=False):
        sta = time.time()
        self._days += 1

        # 记录三阶段人员位置分布
        area_res = []
        area_iso = [False for i in range(self._total_area)]
        # 检查点用
        area_trip = [0 for i in range(self._total_area)]
        err = False

        # 健康状态变化
        for i in range(3):
            self._increase[i] = 0
        for per in self.p_health + self.p_hide + self.p_patient + self.p_no_sym:
            before = per.getStatus()
            cure_rate = self._area_cure['else']
            dead_rate = self._area_dead['else']
            if per.getPosition() == self._hospital_id:
                cure_rate = self._area_cure['hospital']
                dead_rate = self._area_dead['hospital']
            elif per.getPosition() == self._isolation_id:
                cure_rate = self._area_cure['isolation']
                dead_rate = self._area_dead['isolation']
            after = per.defender(cure_rate, dead_rate)
            if before != after:
                self.statusChange(per, before, after)
                if after == 3:
                    self._increase[1] += 1
                elif after == 0 and before == 3:
                    self._increase[0] += 1
                elif after == 4:
                    self._increase[2] += 1
        for per in self.p_dead:
            self.morgue(per)

        # 晚上全回宿舍
        for per in self.p_health:  # 健康人先腾医院和隔离点位置
            self.move(per, per.getDorm())
        # 隔离点确诊优先进医院
        for per in self.p_patient:
            if per.getPosition() == self._isolation_id:
                if self._areas[self._hospital_id].getRemain() > 0:
                    self.move(per, self._hospital_id)
                else:
                    break
        # 确诊进医院或隔离点
        for per in self.p_patient:
            if per.getIsolate() == 0:
                self.move(per, per.getDorm())
                area_iso[per.getPosition()] = True
            elif per.getPosition() != self._isolation_id and per.getPosition() != self._hospital_id:
                if self._areas[self._hospital_id].getRemain() > 0:
                    self.move(per, self._hospital_id)
                elif self._areas[self._isolation_id].getRemain() > 0:
                    self.move(per, self._isolation_id)
                else:
                    self.move(per, per.getDorm())
                    area_iso[per.getPosition()] = True
        # 潜伏或无症状若在隔离点则不动
        for per in self.p_hide + self.p_no_sym:
            if per.getPosition() != self._isolation_id:
                self.move(per, per.getDorm())
        # 设置隔离
        area_iso[self._outdoor_id] = False
        for per in self.p_health + self.p_hide + self.p_no_sym:
            if area_iso[per.getPosition()]:
                per.setIsolate()
        # 完成人员流动后开始进行传染
        self.infect()
        area_res.append(self.areasHandler())
        # 检查点
        if _check:
            print('====checkpoint====')
            for p in self.p_health + self.p_hide + self.p_no_sym + self.p_patient:
                area_trip[p.getPosition()] += 1
            if sum(area_trip) != (self._total_person - len(self.p_dead)):
                err = True
            print(area_trip)

        # 每7天进行一次核酸检测
        if self._confront and self._days % 7 == 0:
            self.detection()
            area_res.append(self.areasHandler())
            area_res.append(self.areasHandler())
            if _check:
                area_trip = [0 for i in range(self._total_area)]
                for p in self.p_health + self.p_hide + self.p_no_sym + self.p_patient:
                    area_trip[p.getPosition()] += 1
                if sum(area_trip) != (self._total_person - len(self.p_dead)):
                    err = True
                print(area_trip)
            return area_res

        # 上午下午随机位置
        for i in range(2):
            area_iso = [False for i in range(self._total_area)]
            for per in self.p_hide + self.p_patient + self.p_no_sym + self.p_health:
                if self._confront and per.getIsolate() > 0:
                    continue
                elif per.getPosition() != self._hospital_id and per.getPosition() != self._isolation_id:
                    # 随机进入图书馆或教室
                    mov = random.randint(self._library_id, self._classroom_id + self._area_num['classroom'] - 1)
                    if self._areas[mov].getRemain() <= 0 or random.random() < 0.1:
                        mov = self._outdoor_id
                    self.move(per, mov)
                    # 病发
                    if per.getStatus() == 3:
                        area_iso[per.getPosition()] = True
            area_iso[self._outdoor_id] = False
            for per in self.p_health + self.p_hide + self.p_no_sym:
                if area_iso[per.getPosition()]:
                    per.setIsolate()
            self.infect()  # 传染
            area_res.append(self.areasHandler())
            # 检查点
            if _check:
                area_trip = [0 for i in range(self._total_area)]
                for p in self.p_health + self.p_hide + self.p_no_sym + self.p_patient:
                    area_trip[p.getPosition()] += 1
                if sum(area_trip) != (self._total_person - len(self.p_dead)):
                    err = True
                print(area_trip)

        # 检查点
        if _check:
            if err:
                print('error!')
            else:
                print('day {} checked'.format(self._days))
            print('==================')
        tt = time.time() - sta
        print('# day{} use time:{}'.format(self._days, tt))
        return area_res

    # 进行传染
    def infect(self):
        # 设置区域风险
        area_virus = [0 for i in range(self._total_area)]
        area_risk = [0 for i in range(self._total_area)]
        for p in self.p_patient:
            area_virus[p.getPosition()] = max(p.getVirus(), area_virus[p.getPosition()])
            area_risk[p.getPosition()] = area_virus[p.getPosition()]
        for p in self.p_no_sym:
            area_virus[p.getPosition()] = max(p.getVirus(), area_virus[p.getPosition()])
            area_risk[p.getPosition()] = max(p.getVirus() * 0.5, area_risk[p.getPosition()])  # 无症状感染率减半
        # 遍历人群进行感染
        for per in self.p_health:
            if random.random() < self._areas[per.getPosition()].getRisk() * area_risk[per.getPosition()]:
                flag = per.beInfected(area_virus[per.getPosition()])
                if flag:
                    self.p_hide.append(per)
                    self.p_health.remove(per)

    # 核酸检测
    def detection(self):
        flag = True
        for per in self.p_patient:
            per.setIsolate()
            if per.getPosition() != self._hospital_id:
                if self._areas[self._hospital_id].getRemain() > 0:
                    self.move(per, self._hospital_id)
                elif self._areas[self._isolation_id].getRemain() > 0:
                    self.move(per, self._isolation_id)
                else:
                    self.move(per, per.getDorm())
                    flag = False
        for per in self.p_no_sym + self.p_hide:
            # 模拟未检出
            if per.getStatus == 1 and random.random() > 0.9:
                continue
            elif per.getStatus == 2 and random.random() > 0.99:
                continue
            # 去隔离
            per.setIsolate()
            if per.getPosition() != self._isolation_id:
                if self._areas[self._isolation_id].getRemain() > 0:
                    self.move(per, self._isolation_id)
                else:
                    self.move(per, per.getDorm())
                    flag = False
        return flag

    # 全员戴口罩   rate有效率
    def allMask(self, rate):
        for per in self.p_health + self.p_hide + self.p_no_sym + self.p_patient:
            per.setProtect(rate)

    # 打疫苗   rate有效率 prop比例
    def vaccinate(self, rate, prop):
        num = self._total_person * prop
        thre = num
        # 应该要根据疫苗间隔时间排序  Person._vaccine_interval 初始值1 +1+1
        hist = [0] * (self._days + 2)
        for per in self.p_hide + self.p_health:
            hist[per.getVaccineInterval()] += 1
        # 反向遍历从大至小
        for i in range(len(hist) - 1, -1, -1):
            thre -= hist[i]
            if thre < 0:
                thre = i
                break
        for per in self.p_hide + self.p_health:
            if per.getVaccineInterval() >= thre:
                per.vaccinate(rate)
                num -= 0
                if num <= 0:
                    break
        # 疫苗全部打完
        if num <= 0:
            return True
        return False

    # 返回每个区域，对应各种状态人的数目
    def areasHandler(self):
        # [[area_id, area_type, 健康数, 潜伏数, 无症状, 病发], ...]
        res = []
        for area in self._areas:
            # [area_id, area_type, 健康数, 潜伏数, 无症状, 病发]
            res.append([area.getID(), area.getType(), 0, 0, 0, 0])
        for per in self.p_health + self.p_hide + self.p_no_sym + self.p_patient:
            # (0: 未感染, 1: 潜伏, 2: 无症状, 3:病发)
            res[per.getPosition()][per.getStatus() + 2] += 1
        return res

    # 异常检测
    def errorCheck(self):
        assert len(self.p_health) + len(self.p_hide) + len(self.p_patient) + len(self.p_no_sym) + len(
            self.p_dead) == self._total_person, '总人数错误'
        flag = True
        for per in self.p_health:
            if per.getStatus() != 0:
                assert flag, '健康队列错误'
        for per in self.p_hide:
            if per.getStatus() != 2:
                assert flag, '潜伏队列错误'
        for per in self.p_no_sym:
            if per.getStatus() != 3:
                assert flag, '无症状队列错误'
        for per in self.p_patient:
            if per.getStatus() != 4:
                assert flag, '确诊队列错误'
        for per in self.p_no_sym:
            if per.getStatus() != 5:
                assert flag, '死亡队列错误'
